# -*- coding: utf-8 -*-
from IfaceMap import *
import ipaddr
from time import *
from scapy.all import get_if_hwaddr, get_if_addr, conf, load_module, Dot11, Dot11Elt
from scapy.error import Scapy_Exception
load_module("p0f")
from scapy.modules.p0f import *

def isLocal(addr):
  # Dirty way to find internal and external IP. 
  if (ipaddr.IP(addr) in ipaddr.IP('10.0.0.0/8')) or (ipaddr.IP(addr) in ipaddr.IP('172.16.0.0/12')) or (ipaddr.IP(addr) in ipaddr.IP('192.168.0.0/16')):
    return True
  return False


class PktHandler:

  def __init__(self,iface):
    
    self.ifmap=IfaceMap(iface)
    self.iface=iface

    try:
      self.mac=get_if_hwaddr(iface)
      self.IP=get_if_addr(iface)
    except Scapy_Exception, msg:
      print '[!] No local address:', msg
      self.mac=None
      self.IP=None
    else:
      if self.mac and not self.__isInvalid(self.mac):
	self.ifmap.set_myMac(self.mac)

      if self.IP and not self.__isInvalid(self.IP):
	self.ifmap.set_myIP(self.IP)

      self.ifmap.set_myNetmask(self.__get_if_addr_slash_netmask(iface))
 
    self.__loadServices()
    self.__loadHwManifacturer()

  def __loadServices(self):
    
    import scapy.data
    
    #NMAP-services. Too entries, it creates too much false-positive open services
    #scapy.data.TCP_SERVICES,scapy.data.UDP_SERVICES=scapy.data.load_services("utils/nmap-services")

    
    self.TCP_SERVICES={}
    self.UDP_SERVICES={}
    
    for p in scapy.data.TCP_SERVICES.keys():
      self.TCP_SERVICES[scapy.data.TCP_SERVICES[p]] = p 
    for p in scapy.data.UDP_SERVICES.keys():
      self.UDP_SERVICES[scapy.data.UDP_SERVICES[p]] = p   

  def __loadHwManifacturer(self):
    
    self.hwList={}
    
    f1 = open("utils/oui.txt","r")
    for line in f1:
        if "(hex)" in line:
	  mac=line.split("(hex)")[0].strip().replace("-",":").lower()
	  manif=line.split("(hex)")[1].strip()
	  self.hwList[mac]=manif
    f1.close()
    
  def __isInvalid(self,addr):
    # Dirty way to find invalid addresses 
    if addr.endswith('.255') or addr=='ff:ff:ff:ff:ff:ff' or addr=='::' or addr.startswith('127.') or addr=='0.0.0.0' or addr=='00:00:00:00:00:00':
      #broadcast and loopback hw, ipv4, ipv6 addresses
      return True
    return False

  #def __isLocal(self,addr):
    ## Dirty way to find internal and external IP. 
    #if (ipaddr.IP(addr) in ipaddr.IP('10.0.0.0/8')) or (ipaddr.IP(addr) in ipaddr.IP('172.16.0.0/12')) or (ipaddr.IP(addr) in ipaddr.IP('192.168.0.0/16')):
      #return True
    #return False

  def __isInNetmask(self,addr,netmask):
    return (ipaddr.IP(addr) in ipaddr.IP(netmask))

  def __get_if_addr_slash_netmask(self,interface):

    for net,msk,gw,iface,addr in conf.route.routes:
      if iface==interface and net!=0:
	
	ip=ipaddr.IP(net).ip_ext + "/" + ipaddr.IP(msk).ip_ext
	return ipaddr.IP(ip)
	

  def writeout(self, path):
    f=open(path,'w+')
    self.ifmap.export(f, 0)
    f.close()


  def procPkt(self,pkt):

      time=strftime("%d-%m-%y %H:%M:%S", localtime())

      # Prepare locale packet variable # <- qua bisognerebbe fare un bel dizionario
      #srcOutLan=False
      #dstOutLan=False
      
      hwSrc=''
      hwDst=''
      ipSrc=''
      ipDst=''
      portSrc=0
      portDst=0
      L3proto=''
      isGw=False
      isMe=False

      # reset old Host object references #
      oldSrc=None
      hSrc=None
      oldDst=None
      hDst=None
      
      #Address 1 is the receiver, Address 2 is the transmitter, 
      #Address 3 is used for filtering purposes by the receiver
      
      if 'Ether' in pkt or 'Dot11' in pkt:
	######################################
	#### Ethernet - Get hSrc e hDest  ####
	######################################
	if 'Ether' in pkt:
	  hwSrc=pkt['Ethernet'].src
	if 'Dot11' in pkt: 
	  hwSrc=pkt[Dot11].addr2

	if(not hwSrc or self.__isInvalid(hwSrc)):
	  return

	for m in self.ifmap.Host:
	  if m.get_MacAddr() == hwSrc:
	    oldSrc = m
	    hSrc = m
	    hSrc.set_LastSeen(time);
	if not hSrc:
	  hSrc = Host(MacAddr=hwSrc,LastSeen=time)
	  if self.mac and hwSrc==self.mac:
	    hSrc.set_Me(True);
	    print '[+] local host detected with MAC', hwSrc
	  
	  if hwSrc[:8] in self.hwList:
	    hSrc.set_MacManifacturer(self.hwList[hwSrc[:8]])
	  else:
	    print "[!]", hwSrc[:8], " address prefix not found in manifacturier list"
	  
	  
	  if 'IP' not in pkt:
	    print '[+] new MAC', hwSrc
	    #pkt.show()
	  if 'Dot11Elt' in pkt:
	      #if not pkt[Dot11Elt].info:
		#pkt[Dot11Elt].info='hidden';
	      #dot11Src = Dot11()
	      #dot11Src.set_Info(pkt[Dot11Elt].info);
	      
	      if len(pkt[Dot11Elt].info.strip())==0:
		info='<hidden>'
	      else:
		info =  pkt[Dot11Elt].info
	      
	      dot11Src=Dot11Info(Info=info)
	      #dot11Src.set_AP(True)
	      hSrc.set_Dot11Info(dot11Src)
	      print '[+] New AP with essid', info, 'and MAC', hwDst
	    
	#DST
	if 'Ether' in pkt:
	  hwDst=pkt['Ethernet'].dst
	if 'Dot11' in pkt:
	  hwDst=pkt[Dot11].addr1
	  
	######################
	##### IP - Add IP ####
	######################
      

	if 'IP' in pkt:
	  ##SRC
	  ipSrc=pkt['IP'].src
	  if(self.__isInvalid(ipSrc)):
	    return
	  if hSrc and isLocal(ipSrc):
	    if ipSrc not in hSrc.get_IPAddr():
	      hSrc.add_IPAddr(ipSrc)
	      if self.IP: 
		hSrc.set_IsInNetmask(self.__isInNetmask(ipSrc,self.__get_if_addr_slash_netmask(self.iface)))
	      print '[+] new IP', ipSrc, 'with MAC', hwSrc
	  else:
	    # ipSrc is not a local address, probably this host is a gw.
	    isGw=True
	    if not hSrc.get_Gw():
	      print '[!]',  hwSrc, 'is a gateway' 
	      hSrc.set_Gw(True)
	   
	   
	  ##DST
	  ipDst=pkt['IP'].dst
	  
	  ##################################    
	  ##### TCP - Add Connections ######
	  ##################################
	  if 'TCP' in pkt:
	    
	    tcpflags=pkt.sprintf("{TCP:%TCP.flags%}")
	    
	    L3Proto='TCP'
	    portSrc=pkt['TCP'].sport
	    portDst=pkt['TCP'].dport
	      
	    connlist = self.ifmap.get_Connection()
	    conn=None
	    source=0
	    
	    # Search if conn already exists yet, and mark who is the Src (IP1 or IP2)
	    for conn in connlist:
	      if conn.get_Proto() == L3Proto:
		if ipSrc == conn.get_IP1() and portSrc == conn.get_Port1() and ipDst==conn.get_IP2() and portDst==conn.get_Port2():
		  oldconn=conn
		  source=1
		  break
		if ipSrc == conn.get_IP2() and portSrc == conn.get_Port2() and ipDst==conn.get_IP1() and portDst==conn.get_Port1():
		  oldconn=conn
		  source=2
		  break
	    
	    
	    #######################################
	    # OS fingerprinting with SYN packets  #
	    #######################################
	    if 'S' in tcpflags and not ('A' in tcpflags or 'R' in tcpflags or 'F' in tcpflags):
	      try:
		oslist=p0f(pkt)
	      except:
		print '[!] OS fingerprinting failed.' 
	      if len(oslist)>0:
		os = oslist[0]
		oldos = hSrc.get_OperatingSystem()
		if not oldos or oldos.get_Name() != os[0] or oldos.get_Version() != os[1]:
		  hSrc.set_OperatingSystem(OperatingSystem(Name=os[0], Version=os[1]))
		  print '[+] operating system', os[0], os[1], 'running on', ipSrc
		  
	    #############################
	    # Already know connections	#
	    ############################
	    # Posso affinare il 3WH guardando se parte dal server o dal client

	    if conn and source>0:
 
	      status=conn.get_Status()
	      
	      ## controlla se è locale
	      if isLocal(ipSrc) and isLocal(ipDst):
		conn.set_toOutside(False)
	      else:
		conn.set_toOutside(True)
	      
	      # SOMEONE SEND SYN
	      #SYN (ACK, RST, and FIN off)
	      if 'S' in tcpflags and not ('A' in tcpflags or 'R' in tcpflags or 'F' in tcpflags):
		  conn.set_Status('3WH1')
		  
	      # SOMEONE SEND SYN+ACK
	      #SYN/ACK (RST and FIN off)
	      elif 'S' in tcpflags and 'A' in tcpflags and not ('R' in tcpflags or 'F' in tcpflags):
		if(status=='3WH1'):
		  conn.set_Status('3WH2')
		  # Teoricamente da fare OS detection, anche se mi dice che non c'e' SYN. magari provare con tcpflags
		  
	      # SOMEONE SEND ACK
	      #ACK (RST, SYN, and FIN off)
	      elif 'A' in tcpflags and not ('S' in tcpflags or 'R' in tcpflags or 'F' in tcpflags):
		if(status=='3WH2'):
		  conn.set_Status('OPENED')
		  #print '[+] new', conn.get_IP1() + ":" + str(portSrc), '->', conn.get_IP2() + ":" + str(portDst), '[' + conn.get_ServiceName() + ']'
	

		if status=='HCLOSED':
		  conn.set_Status('CLOSED')
		  #connlist.remove(oldconn)
		  #self.ifmap.set_Connection(connlist)
	      
	    
	      #FIN (SYN/RST off)
	      elif 'F' in tcpflags and not ('S' in tcpflags or 'R' in tcpflags):
		if status=='HCLOSED':
		  conn.set_Status('CLOSED')
		  #connlist.remove(oldconn)
		  #self.ifmap.set_Connection(connlist)
		else:
		  conn.set_Status('HCLOSED')

		  
	      #RST (SYN and FIN off)
	      elif 'R' in tcpflags and not ('S' in tcpflags or 'F' in tcpflags):
		if(status=='3WH1'):
		  #se mi avevano precedentemente mandato un ACK, scarto direttamente il pacchetto (non aggiungo nnte).
		  return
		conn.set_Status('CLOSED')
		
	      
	      if oldconn != conn:
		connlist.remove(oldconn)
		connlist.append(conn)
		self.ifmap.set_Connection(connlist)
	      
	    ###################
	    ## New Connection #
	    ###################
	    if source==0:

	      newconn=Connection(Proto=L3Proto, \
	      MAC1=hwSrc, IP1=ipSrc, Port1=portSrc,\
	      MAC2=hwDst, IP2=ipDst, Port2=portDst,\
	      LastSeen=time)
	  
	      ## controlla se è locale
	      if isLocal(ipSrc) and isLocal(ipDst):
		newconn.set_toOutside(False)
	      else:
		newconn.set_toOutside(True)

	      if portSrc in self.TCP_SERVICES:
		servicename=self.TCP_SERVICES[portSrc]
		newconn.set_ServerIs(1)
	      elif portDst in self.TCP_SERVICES:
		servicename=self.TCP_SERVICES[portDst]
		newconn.set_ServerIs(2)
	      else:
		servicename='Unknown'
	      newconn.set_ServiceName(servicename)
	      
	      #Una connessione appena instaurata, può essere solo un SYN o un flusso dati non riconosciuto
	      
	      #SYN (ACK, RST, and FIN off)
	      if 'S' in tcpflags and not ('A' in tcpflags or 'R' in tcpflags or 'F' in tcpflags):
		newconn.set_Status('3WH1')
		
	      #RST or FIN - if the first packet is a TCP or RST, drop packet   
	      elif 'R' in tcpflags or 'F' in tcpflags:
		return
	      else:
		newconn.set_Status('ACTIVE')
		
	      self.ifmap.add_Connection(newconn)
	      
	      
	  ###################################    
	  ##### TCP/UDP - Add Services ###### 
	  ###################################
	    if isLocal(ipSrc):
	      ############ SRC Services ########### 
	      servlist = hSrc.get_Service()
	      thereis=False
	      for serv in servlist:
		if serv.get_Port() == portSrc:
		  thereis=True
		  
	      if thereis==False:
		if portSrc in self.TCP_SERVICES:
		  # It's considered as a service only if in TCP_SERVICES
		  name=self.TCP_SERVICES[portSrc]
		  hSrc.add_Service(Service(Name=name, Port=portSrc, Proto=L3Proto))
		  print '[+] new service ' + ipSrc + ':' + str(portSrc), '[' + name + '/' + L3Proto + ']'
	  
	  
	  ###################################    
	  ##### UDP - Add Services ###### 
	  ###################################
	  
	  if 'UDP' in pkt:
	    
	    L3Proto='UDP'
	    portSrc=pkt['UDP'].sport
	    portDst=pkt['UDP'].dport
	    
	    if isLocal(ipSrc):
	      ############ SRC Services ########### 
	      servlist = hSrc.get_Service()
	      thereis=False
	      for serv in servlist:
		if serv.get_Port() == portSrc:
		  thereis=True
		  
	      if thereis==False:
		if portSrc in self.UDP_SERVICES:
		  # It's considered as a service only if in TCP_SERVICES
		  name=self.UDP_SERVICES[portSrc]
		  hSrc.add_Service(Service(Name=name, Port=portSrc, Proto=L3Proto))
		  print '[+] new service ' + ipSrc + ':' + str(portSrc), '[' + name + '/' + L3Proto + ']'


      ### Save new Host in self.ifmap.Host
      
      if hSrc:
	if oldSrc and (oldSrc in self.ifmap.Host):
	  self.ifmap.Host.remove(oldSrc)
	  self.ifmap.Host.append(hSrc)
	else:
	  self.ifmap.Host.append(hSrc)
      
      
   
